"""
配对交易策略(pyfolio实现)
策略逻辑：
1. 选择两个协整的资产
2. 计算价差的Z-Score
3. 当价差异常时做多低估资产，做空高估资产
"""

import pandas as pd
import numpy as np
import statsmodels.api as sm
import pyfolio as pf

def pairs_trading_strategy(prices_A, prices_B, lookback=30, zscore_threshold=2.0):
    """
    配对交易策略
    
    参数:
        prices_A: 资产A价格序列
        prices_B: 资产B价格序列
        lookback: 计算窗口
        zscore_threshold: 交易阈值
    
    返回:
        positions: 持仓DataFrame(包含asset_A和asset_B两列)
    """
    # 合并价格数据
    df = pd.DataFrame({
        'A': prices_A,
        'B': prices_B
    }).dropna()
    
    # 初始化持仓
    positions = pd.DataFrame(0, index=df.index, columns=['A', 'B'])
    
    for i in range(lookback, len(df)):
        # 计算价差
        spread = df['A'].iloc[i-lookback:i] - df['B'].iloc[i-lookback:i]
        
        # 计算Z-Score
        mean = spread.mean()
        std = spread.std()
        zscore = (spread.iloc[-1] - mean) / std
        
        # 交易信号
        if zscore > zscore_threshold:
            positions.iloc[i] = [-1, 1]  # 做空A，做多B
        elif zscore < -zscore_threshold:
            positions.iloc[i] = [1, -1]  # 做多A，做空B
        else:
            positions.iloc[i] = [0, 0]  # 平仓
    
    return positions

if __name__ == '__main__':
    # 示例用法
    import yfinance as yf
    
    # 获取示例数据(如可口可乐和百事可乐)
    ko = yf.download('KO', start='2020-01-01', end='2023-01-01')['Close']
    pep = yf.download('PEP', start='2020-01-01', end='2023-01-01')['Close']
    
    # 运行策略
    positions = pairs_trading_strategy(ko, pep)
    
    # 计算组合收益
    returns = positions.shift(1) * pd.DataFrame({
        'A': ko.pct_change(),
        'B': pep.pct_change()
    })
    pf.create_returns_tear_sheet(returns.sum(axis=1))